from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import requests
import json
import os
import tempfile
from pathlib import Path
from typing import List, Dict, Any, Optional
import uvicorn
import pandas as pd
from io import BytesIO
import urllib.parse
import logging
import traceback
import re
from text_highlighter import highlight_text_in_content, highlight_multiple_terms_in_content

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="文档整理API", description="AI文档整理服务")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置
API_KEY = ""
BASE_URL = "https://geekai.co/api/v1/chat/completions"


# 请求模型定义
class MultiTableRequest(BaseModel):
    content: str
    user_request: str
    messages_history: List[Dict[str, str]] = []


class ExcelExportRequest(BaseModel):
    data: Optional[List[Dict[str, Any]]] = None
    tables: Optional[List[Dict[str, Any]]] = None
    filename: str = "文档数据"
    sheet_name: str = "Sheet1"
    type: str = "single_sheet"


class HighlightRequest(BaseModel):
    original_text: str
    search_text: str


class MultipleHighlightRequest(BaseModel):
    original_text: str
    search_terms: List[str]


class DocumentProcessor:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }

    def call_ai_with_history(self, messages: List[Dict[str, str]]) -> str:
        """调用AI API with完整对话历史"""
        data = {
            "model": "gpt-5-mini",
            "messages": messages
        }

        try:
            response = requests.post(BASE_URL, headers=self.headers, json=data)
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"AI API调用失败: {str(e)}")

    def call_ai(self, content: str) -> str:
        """简单AI调用"""
        data = {
            "model": "gpt-5-mini",
            "messages": [{"role": "user", "content": content}]
        }

        try:
            response = requests.post(BASE_URL, headers=self.headers, json=data)
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"AI API调用失败: {str(e)}")

    def read_file_content(self, file_path: str) -> str:
        """读取文件内容"""
        file_path = Path(file_path)
        file_extension = file_path.suffix.lower()

        try:
            if file_extension == '.txt':
                return self._read_txt(file_path)
            elif file_extension == '.docx':
                return self._read_docx(file_path)
            elif file_extension == '.pdf':
                return self._read_pdf(file_path)
            elif file_extension in ['.md', '.markdown']:
                return self._read_txt(file_path)
            else:
                return self._read_txt(file_path)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"文件读取失败: {str(e)}")

    def _read_txt(self, file_path: Path) -> str:
        encodings = ['utf-8', 'gbk', 'gb2312', 'latin1']
        for encoding in encodings:
            try:
                with open(file_path, 'r', encoding=encoding) as file:
                    return file.read()
            except UnicodeDecodeError:
                continue
        raise Exception("无法识别文件编码")

    def _read_docx(self, file_path: Path) -> str:
        try:
            from docx import Document
            doc = Document(file_path)
            content = []
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    content.append(paragraph.text)
            for table in doc.tables:
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        if cell.text.strip():
                            row_text.append(cell.text.strip())
                    if row_text:
                        content.append(" | ".join(row_text))
            return "\n".join(content)
        except ImportError:
            raise Exception("需要安装 python-docx 库")

    def _read_pdf(self, file_path: Path) -> str:
        try:
            import PyPDF2
            content = []
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                for page in pdf_reader.pages:
                    text = page.extract_text()
                    if text.strip():
                        content.append(text)
            return "\n".join(content)
        except ImportError:
            raise Exception("需要安装 PyPDF2 库")

    def process_file_initial(self, content: str) -> dict:
        """文件首次上传处理"""

        prompt = f"""分析以下文档内容，智能生成合适的表格。

文档内容：
{content}

严格按照以下JSON格式返回，不允许有任何其他内容：

{{"tables":[{{"name":"表格名称","description":"表格用途","data":[{{"列名1":"数据1","列名2":"数据2"}}]}}],"summary":"处理说明"}}

要求：
1. 必须是一行完整的JSON，不能换行
2. 所有字符串必须用双引号
3. 不能有任何JSON之外的文字、解释或格式
4. 返回内容的第一个字符必须是{{，最后一个字符必须是}}
5. 根据文档内容智能决定生成几个表格和列结构

示例输出：
{{"tables":[{{"name":"学习内容表","description":"课程内容整理","data":[{{"序号":"1","标题":"Unit 1","内容":"Hello world"}},{{"序号":"2","标题":"Unit 2","内容":"Good morning"}}]}}],"summary":"生成了1个学习内容表，包含2条记录"}}"""

        try:
            ai_response = self.call_ai(prompt).strip()
            logger.info(f"AI初始响应: {ai_response[:200]}...")

            # 直接解析
            result = json.loads(ai_response)

            return {
                "success": True,
                "tables": result.get("tables", []),
                "summary": result.get("summary", "处理完成"),
                "processing_type": "multiple_tables"
            }

        except json.JSONDecodeError as e:
            logger.error(f"JSON解析失败: {str(e)}")
            logger.error(f"AI响应内容: {ai_response}")
            # 兜底处理
            return {
                "success": True,
                "tables": [{
                    "name": "文档内容表",
                    "description": "文档内容整理",
                    "data": [{"序号": "1", "内容": content[:1000] + "..." if len(content) > 1000 else content}]
                }],
                "summary": "AI解析失败，使用默认格式",
                "processing_type": "multiple_tables"
            }
        except Exception as e:
            logger.error(f"处理失败: {str(e)}")
            return {
                "success": True,
                "tables": [{
                    "name": "文档内容表",
                    "description": "文档内容整理",
                    "data": [{"序号": "1", "内容": content[:1000] + "..." if len(content) > 1000 else content}]
                }],
                "summary": f"处理异常: {str(e)}",
                "processing_type": "multiple_tables"
            }

    def generate_tables_with_history(self, content: str, user_request: str, messages_history: List[Dict]) -> dict:
        """基于完整对话历史生成表格"""

        messages = [
            {
                "role": "system",
                "content": f"""你是专业的文档分析助手。文档内容：

{content}

严格按照以下JSON格式返回，不允许有任何其他内容：

{{"tables":[{{"name":"表格名称","description":"表格用途","data":[{{"列名1":"数据1","列名2":"数据2"}}]}}],"summary":"处理说明"}}

要求：
1. 必须是一行完整的JSON，不能换行
2. 所有字符串必须用双引号
3. 不能有任何JSON之外的文字、解释或格式
4. 如果数据中有引号，必须转义为\"
5. 返回内容的第一个字符必须是{{，最后一个字符必须是}}

示例输出：
{{"tables":[{{"name":"词汇表","description":"单词整理","data":[{{"英文":"hello","中文":"你好"}},{{"英文":"world","中文":"世界"}}]}}],"summary":"生成了1个词汇表，包含2个单词"}}"""
            }
        ]

        messages.extend(messages_history)
        messages.append({
            "role": "user",
            "content": f"用户需求：{user_request}\n\n严格按照上述JSON格式返回，不要任何其他内容。"
        })

        try:
            ai_response = self.call_ai_with_history(messages).strip()
            logger.info(f"AI对话响应: {ai_response[:200]}...")

            # 直接解析
            result = json.loads(ai_response)

            return {
                "tables": result.get("tables", []),
                "summary": result.get("summary", "处理完成")
            }

        except json.JSONDecodeError as e:
            logger.error(f"JSON解析失败: {str(e)}")
            logger.error(f"AI响应内容: {ai_response}")
            # 兜底处理
            return {
                "tables": [{
                    "name": "处理失败表",
                    "description": f"根据需求'{user_request}'处理失败",
                    "data": [{"错误信息": str(e), "用户需求": user_request}]
                }],
                "summary": "AI响应格式错误，请重试"
            }
        except Exception as e:
            logger.error(f"生成表格失败: {str(e)}")
            return {
                "tables": [{
                    "name": "错误表",
                    "description": f"处理用户需求时出错",
                    "data": [{"错误": str(e), "需求": user_request}]
                }],
                "summary": f"处理失败: {str(e)}"
            }

    def export_to_excel(self, data: Optional[List[Dict[str, Any]]] = None,
                        tables: Optional[List[Dict[str, Any]]] = None,
                        filename: str = "数据导出",
                        sheet_name: str = "Sheet1",
                        export_type: str = "single_sheet") -> BytesIO:
        """导出数据到Excel"""
        try:
            output = BytesIO()

            if export_type == "multiple_sheets" and tables:
                with pd.ExcelWriter(output, engine='openpyxl') as writer:
                    for table_info in tables:
                        table_data = table_info.get('data', [])
                        table_name = table_info.get('name', 'Data')

                        if not table_data:
                            continue

                        safe_sheet_name = table_name.replace('/', '_').replace('\\', '_')[:31]
                        cleaned_data = self._clean_data_for_excel(table_data)

                        if cleaned_data:
                            df = pd.DataFrame(cleaned_data)
                            df.to_excel(writer, sheet_name=safe_sheet_name, index=False)

                            worksheet = writer.sheets[safe_sheet_name]
                            for column in df.columns:
                                column_length = max(
                                    df[column].astype(str).apply(len).max(),
                                    len(str(column))
                                )
                                column_length = min(max(column_length, 10), 50)
                                col_letter = worksheet.cell(row=1, column=df.columns.get_loc(column) + 1).column_letter
                                worksheet.column_dimensions[col_letter].width = column_length
            else:
                if not data:
                    raise ValueError("没有数据可导出")

                cleaned_data = self._clean_data_for_excel(data)
                df = pd.DataFrame(cleaned_data)

                with pd.ExcelWriter(output, engine='openpyxl') as writer:
                    df.to_excel(writer, sheet_name=sheet_name, index=False)
                    worksheet = writer.sheets[sheet_name]

                    for column in df.columns:
                        column_length = max(
                            df[column].astype(str).apply(len).max(),
                            len(str(column))
                        )
                        column_length = min(max(column_length, 10), 50)
                        col_letter = worksheet.cell(row=1, column=df.columns.get_loc(column) + 1).column_letter
                        worksheet.column_dimensions[col_letter].width = column_length

            output.seek(0)
            return output

        except Exception as e:
            logger.error(f"Excel导出失败: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Excel导出失败: {str(e)}")

    def _clean_data_for_excel(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """清理数据确保Excel兼容性"""
        cleaned_data = []
        for row in data:
            cleaned_row = {}
            for key, value in row.items():
                clean_key = str(key).strip() or "未命名字段"
                if value is None:
                    clean_value = ""
                else:
                    clean_value = str(value).strip()
                    if len(clean_value) > 32767:
                        clean_value = clean_value[:32764] + "..."
                cleaned_row[clean_key] = clean_value
            cleaned_data.append(cleaned_row)
        return cleaned_data


# 全局处理器
processor = None


@app.on_event("startup")
async def startup_event():
    global processor
    api_key = os.getenv("GEEKAI_API_KEY")
    if not api_key:
        logger.warning("未设置GEEKAI_API_KEY环境变量")
    processor = DocumentProcessor(api_key)


@app.post("/api/upload-file")
async def upload_file(file: UploadFile = File(...)):
    """上传并处理文件"""
    global processor
    if not processor:
        raise HTTPException(status_code=500, detail="服务未初始化")

    allowed_extensions = {'.txt', '.docx', '.pdf', '.md', '.markdown'}
    file_extension = Path(file.filename).suffix.lower()

    if file_extension not in allowed_extensions:
        raise HTTPException(
            status_code=400,
            detail=f"不支持的文件类型: {file_extension}。支持的类型: {', '.join(allowed_extensions)}"
        )

    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
            content = await file.read()
            tmp_file.write(content)
            tmp_file_path = tmp_file.name

        try:
            # 读取文件内容
            file_content = processor.read_file_content(tmp_file_path)
            logger.info(f"文件内容读取完成，长度: {len(file_content)}")

            # 智能生成初始表格
            result = processor.process_file_initial(file_content)

            return {
                "success": True,
                "message": "文件处理成功",
                "filename": file.filename,
                "size": len(content),
                "processing_type": result["processing_type"],
                "tables": result["tables"],
                "summary": result["summary"],
                "originalContent": file_content
            }

        finally:
            os.unlink(tmp_file_path)

    except Exception as e:
        logger.error(f"文件处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"文件处理失败: {str(e)}")


@app.post("/api/generate-multiple-tables")
async def generate_multiple_tables_api(request: MultiTableRequest):
    """基于对话历史生成表格"""
    global processor
    if not processor:
        raise HTTPException(status_code=500, detail="服务未初始化")

    try:
        logger.info(f"收到用户需求：'{request.user_request}'")
        logger.info(f"对话历史：{len(request.messages_history)} 条消息")

        result = processor.generate_tables_with_history(
            content=request.content,
            user_request=request.user_request,
            messages_history=request.messages_history
        )

        return {
            "success": True,
            "message": "处理完成",
            "result": result
        }

    except Exception as e:
        logger.error(f"处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")


@app.post("/api/highlight-text")
async def highlight_text_api(request: HighlightRequest):
    """高亮文本接口"""
    try:
        if not request.original_text:
            raise HTTPException(status_code=400, detail="原文不能为空")

        if not request.search_text or not request.search_text.strip():
            from text_highlighter import highlighter
            result = f'<span class="highlight-normal">{highlighter.escape_html(request.original_text).replace(chr(10), "<br>")}</span>'
        else:
            result = highlight_text_in_content(request.original_text, request.search_text.strip())

        return {
            "success": True,
            "message": "文本高亮处理成功",
            "highlighted_text": result,
            "search_text": request.search_text,
            "original_length": len(request.original_text)
        }

    except Exception as e:
        logger.error(f"文本高亮处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"文本高亮处理失败: {str(e)}")


@app.post("/api/highlight-multiple")
async def highlight_multiple_text_api(request: MultipleHighlightRequest):
    """高亮多个搜索词接口"""
    try:
        if not request.original_text:
            raise HTTPException(status_code=400, detail="原文不能为空")

        if not request.search_terms or not any(term.strip() for term in request.search_terms):
            from text_highlighter import highlighter
            result = f'<span class="highlight-normal">{highlighter.escape_html(request.original_text).replace(chr(10), "<br>")}</span>'
        else:
            valid_terms = [term.strip() for term in request.search_terms if term.strip()]
            result = highlight_multiple_terms_in_content(request.original_text, valid_terms)

        return {
            "success": True,
            "message": "多词高亮处理成功",
            "highlighted_text": result,
            "search_terms": request.search_terms,
            "original_length": len(request.original_text)
        }

    except Exception as e:
        logger.error(f"多词高亮处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"多词高亮处理失败: {str(e)}")


@app.post("/api/export-excel")
async def export_excel(request: ExcelExportRequest):
    """导出Excel文件"""
    global processor
    if not processor:
        raise HTTPException(status_code=500, detail="服务未初始化")

    try:
        excel_buffer = processor.export_to_excel(
            data=request.data,
            tables=request.tables,
            filename=request.filename,
            sheet_name=request.sheet_name,
            export_type=request.type
        )

        safe_filename = request.filename.replace(' ', '_').replace('/', '_').replace('\\', '_')
        from datetime import datetime
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        final_filename = f"{safe_filename}_{timestamp}.xlsx"

        encoded_filename = urllib.parse.quote(final_filename.encode('utf-8'))
        content_disposition = f"attachment; filename*=UTF-8''{encoded_filename}"

        return StreamingResponse(
            BytesIO(excel_buffer.read()),
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            headers={
                "Content-Disposition": content_disposition,
                "Access-Control-Expose-Headers": "Content-Disposition"
            }
        )

    except Exception as e:
        logger.error(f"Excel导出失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Excel导出失败: {str(e)}")


@app.get("/api/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy", "message": "服务运行正常"}


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        os.environ["GEEKAI_API_KEY"] = sys.argv[1]
    else:
        api_key = input("请输入GEEKAI API密钥: ")
        os.environ["GEEKAI_API_KEY"] = api_key
    uvicorn.run(app, host="0.0.0.0", port=8801)
